#include "corrections.h"
#include "physical_properties.h"

//classical
void Compute_long_range_corrections(double& long_range_correction_energy, double& long_range_correction_pressure) {
    long_range_correction_energy = 4.0 * Constant_Epsilon * pow(Constant_Sigma, 3) * pow(Parameter_Natoms, 2) * (pow((Constant_Sigma / Parameter_Rcutoff), 9) / 3.0 - pow((Constant_Sigma / Parameter_Rcutoff), 3));
    long_range_correction_energy *= 2.0 * Constant_Pi / (3.0 * pow(Parameter_Box_size, 3));

    long_range_correction_pressure = 4.0 * Constant_Epsilon * pow(Constant_Sigma, 3) * pow(Parameter_Natoms, 2) * (pow((Constant_Sigma / Parameter_Rcutoff), 9) * 4.0 - pow((Constant_Sigma / Parameter_Rcutoff), 3) * 6.0);
    long_range_correction_pressure *= 2.0 * Constant_Pi / (3.0 * pow(Parameter_Box_size, 3));

    long_range_correction_pressure /= (3.0 * pow(Parameter_Box_size, 3));
}

void inside_box(std::vector<Atom>& mol) {
    int Natoms = mol.size();
    for (int i = 0; i < Natoms; i++) {
        for (int j = 0; j < 3; j++) {
            if (mol[i].coord[j] > Parameter_Box_size) {
                mol[i].coord[j] -= Parameter_Box_size;
            }
            if (mol[i].coord[j] < 0.0) {
                mol[i].coord[j] += Parameter_Box_size;
            }
        }
    }
}

void remove_com_velocity(std::vector<Atom>& mol) {
    std::vector<double> velocity_com({ 0, 0, 0 });
    int Natoms = mol.size();
    for (int i = 0; i < Natoms; i++){
        for (int j = 0; j < 3; j++){
            velocity_com[j] += mol[i].veloc[j];
        }
    }

    for (int i = 0; i < 3; i++)
        velocity_com[i] /= double(Natoms);

    for (int i = 0; i < Natoms; i++){
        for (int j = 0; j < 3; j++){
            mol[i].veloc[j] -= velocity_com[j];
        }
    }
    return;
}

void rescale_velocities(std::vector<Atom>& mol){
    double Ek = mol_sys::Compute_sys_Ek(mol);
    double temperature = mol_sys::Compute_Temperature(mol, Ek);

    double scaling_factor_squared = 1.0 + (Parameter_Timestep_fs / Parameter_Thermostat_relaxation_time_fs) * (Parameter_Temperature_imposed / temperature - 1.0);
    double scaling_factor = sqrt(scaling_factor_squared);
    for (int i = 0; i < mol.size(); i++)
        for (int j = 0; j < 3; j++)
            mol[i].veloc[j] *= scaling_factor;

    return;
}

/**************************************************************************PIMD**************************************************************************/
void inside_box_PIMD(std::vector< std::vector<Atom> >& mirr) {
    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < Parameter_Natoms; i++) {
            for (int j = 0; j < 3; j++) {
                if (mirr[k][i].coord[j] > Parameter_Box_size)mirr[k][i].coord[j] -= Parameter_Box_size;
                if (mirr[k][i].coord[j] < 0.0) mirr[k][i].coord[j] += Parameter_Box_size;
            }
        }
    }
}

void remove_com_velocity_PIMD(std::vector< std::vector<Atom> >& mirr) {
    for (int k = 0; k < Parameter_nbeads; k++)
        remove_com_velocity(mirr[k]);
}

void rescale_velocities_PIMD(std::vector< std::vector<Atom> >& mirr) {
    double Ek = mol_sys::Compute_Ek_stage(staging_mirror) / Parameter_nbeads;
    double temperature = mol_sys::Compute_Temperature(mirr[0], Ek);

    double scaling_factor_squared = 1.0 + (Parameter_Timestep_fs / Parameter_Thermostat_relaxation_time_fs) * (Parameter_Temperature_imposed / temperature - 1.0);
    double scaling_factor = sqrt(scaling_factor_squared);

    for (int k = 0; k < Parameter_nbeads; k++) {
        for (int i = 0; i < mirr[k].size(); i++)
            for (int j = 0; j < 3; j++)
                mirr[k][i].veloc[j] *= scaling_factor;
    }
}


